package wang.yang;

import javafx.util.Pair;
import org.apache.commons.collections.CollectionUtils;
import org.apache.zookeeper.*;
import org.apache.zookeeper.data.Stat;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import wang.yang.util.ProcessRunner;

import java.io.File;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.Objects;
import java.util.Random;
import java.util.concurrent.ArrayBlockingQueue;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;

/**
 * 获取已分配的任务
 * 执行并创建状态节点
 */
public class MyWorker implements Watcher {
    private static final Logger LOG = LoggerFactory.getLogger(MyWorker.class);
    private static final byte[] ZERO_LENGTH_ARR = new byte[0];

    private ZooKeeper zk;
    private String hostPort;
    private String serverId = Integer.toHexString((new Random()).nextInt());
    String workerName = "worker-" + serverId;
    String status;
    private ProcessRunner runner = new ProcessRunner();
    private List<String> assignedTaskListCache;
    private volatile boolean connected = false;
    private volatile boolean expired = false;

    /*
     * In general, it is not a good idea to block the callback thread
     * of the ZooKeeper client. We use a thread pool executor to detach
     * the computation from the callback.
     */
    // 创建线程池执行任务
    private ThreadPoolExecutor executor;

    /**
     * Creates a new Worker instance.
     *
     * @param hostPort
     */
    public MyWorker(String hostPort) {
        this.hostPort = hostPort;
        // 只有一个线程，一次只能执行单个任务
        this.executor = new ThreadPoolExecutor(1, 1,
                1000L,
                TimeUnit.MILLISECONDS,
                new ArrayBlockingQueue<Runnable>(200));
    }

    void startZK() throws IOException {
        zk = new ZooKeeper(hostPort, 5000, this);
    }

    @Override
    public void process(WatchedEvent e) {
        LOG.info(e.toString() + ", " + hostPort);
        if (e.getType() == Event.EventType.None) {
            switch (e.getState()) {
                case SyncConnected:
                    connected = true;
                    break;
                case Disconnected:
                    connected = false;
                    break;
                case Expired:
                    expired = true;
                    connected = false;
                    LOG.error("Session expired");
                default:
                    break;
            }
        }
    }

    public void createWorkerAndAssignNode() {
        Transaction transaction = zk.transaction();
        // 创建worker节点 和 任务分配路径
        transaction.create("/workers/" + workerName, "Idle".getBytes(), ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.EPHEMERAL);
        transaction.create("/assign/worker-" + serverId, ZERO_LENGTH_ARR, ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT);
        transaction.commit(createWorkerAndAssignNodeCallback, null);
    }

    AsyncCallback.MultiCallback createWorkerAndAssignNodeCallback = new AsyncCallback.MultiCallback() {
        @Override
        public void processResult(int rc, String path, Object ctx, List<OpResult> opResults) {
            switch (KeeperException.Code.get(rc)) {
                case CONNECTIONLOSS:
                    createWorkerAndAssignNode();
                    break;
                case OK:
                    OpResult.CreateResult workerCreateResult = (OpResult.CreateResult) opResults.get(0);
                    OpResult.CreateResult assignNodeCreateResult = (OpResult.CreateResult) opResults.get(1);
                    LOG.info("Assign node And worker node created: {}, {}", workerCreateResult.getPath(), assignNodeCreateResult.getPath());
                    break;
                case NODEEXISTS:
                    LOG.warn("Assign node and worker node already registered, {}", path);
                    break;
                default:
                    LOG.error("Something went wrong: " + KeeperException.create(KeeperException.Code.get(rc), path));
            }
        }
    };

    AsyncCallback.StatCallback statusUpdateCallback = new AsyncCallback.StatCallback() {
        @Override
        public void processResult(int rc, String path, Object ctx, Stat stat) {
            switch (KeeperException.Code.get(rc)) {
                case CONNECTIONLOSS:
                    updateStatus((String) ctx);
                    return;
            }
        }
    };


    synchronized private void updateStatus(String status) {
        if (status == this.status) {
            zk.setData("/workers/" + workerName, status.getBytes(), -1,
                    statusUpdateCallback, status);
        }
    }

    public void setStatus(String status) {
        this.status = status;
        updateStatus(status);
    }

    /*
     ***************************************
     ***************************************
     * Methods to wait for new assignments.*
     ***************************************
     ***************************************
     */

    Watcher assignedTaskChangeWatcher = new Watcher() {
        @Override
        public void process(WatchedEvent e) {
            if (e.getType() == Event.EventType.NodeChildrenChanged) {
                getAssignedTasks();
            }
        }
    };

    void getAssignedTasks() {
        // 获取已分配的任务
        zk.getChildren("/assign/" + workerName,
                assignedTaskChangeWatcher,
                getAssignedTasksCallback,
                null);
    }

    AsyncCallback.ChildrenCallback getAssignedTasksCallback = new AsyncCallback.ChildrenCallback() {
        @Override
        public void processResult(int rc, String path, Object ctx, List<String> assignedTaskList) {
            switch (KeeperException.Code.get(rc)) {
                case CONNECTIONLOSS:
                    getAssignedTasks();
                    break;
                case OK:
                    if (CollectionUtils.isNotEmpty(assignedTaskList)) {
                        // 先执行获取任务列表操作
                        List<String> newlyAssignedTaskList;
                        if (Objects.isNull(assignedTaskListCache)) {
                            newlyAssignedTaskList = assignedTaskList;
                            assignedTaskListCache = assignedTaskList;
                        } else {
                            newlyAssignedTaskList = new ArrayList<>(assignedTaskList);
                            newlyAssignedTaskList.removeAll(assignedTaskListCache);
                            assignedTaskListCache = assignedTaskList;
                        }
                        if (CollectionUtils.isEmpty(newlyAssignedTaskList)) {
                            return;
                        }

                        LOG.info("获取新分配的任务并执行");
                        setStatus("Working");
                        // 遍历新增的任务
                        for (String taskId : newlyAssignedTaskList) {
                            LOG.info("New task: {}", taskId);
                            // 获取任务信息
                            zk.getData("/assign/" + workerName + "/" + taskId,
                                    false,
                                    getTaskDataCallback,
                                    taskId);
                        }
                    }
                    break;
                default:
                    System.out.println("getChildren failed: " + KeeperException.create(KeeperException.Code.get(rc), path));
            }
        }
    };

    AsyncCallback.DataCallback getTaskDataCallback = new AsyncCallback.DataCallback() {
        @Override
        public void processResult(int rc, String path, Object taskId, byte[] commandData, Stat stat) {
            switch (KeeperException.Code.get(rc)) {
                case CONNECTIONLOSS:
                    zk.getData(path, false, getTaskDataCallback, null);
                    break;
                case OK:
                    /*
                     *  Executing a task in this example is simply printing out
                     *  some string representing the task.
                     */
                    executor.execute(() -> {
                        // 线程池中执行任务(用户目录下执行shell命令)
                        runner.run(new String(commandData), false, new File(System.getProperty("user.home")));
                        String result = String.join("\n", runner.getOutputList());
                        LOG.info("taskId {} 执行结果:{}", (String) taskId, result);
                        // 创建任务状态节点 并删除 已分配的任务
                        createStatusNodeAndDeleteTask((String) taskId, result);
                    });
                    break;
                default:
                    LOG.error("Failed to get task data: ", KeeperException.create(KeeperException.Code.get(rc), path));
            }
        }
    };

    private void createStatusNodeAndDeleteTask(String taskId, String result) {
        Transaction transaction = zk.transaction();
        transaction.create("/status/" + taskId, result.getBytes(), ZooDefs.Ids.OPEN_ACL_UNSAFE,
                CreateMode.PERSISTENT);
        transaction.delete("/assign/" + workerName + "/" + taskId, -1);
        Pair<String, String> taskIdAndResult = new Pair<>(taskId, result);
        transaction.commit(createStatusNodeAndDeleteTaskCallback, taskIdAndResult);
    }

    AsyncCallback.MultiCallback createStatusNodeAndDeleteTaskCallback = new AsyncCallback.MultiCallback() {
        @Override
        public void processResult(int rc, String path, Object taskIdAndResult, List<OpResult> opResults) {
            Pair<String, String> taskIdAndResultPair = (Pair<String, String>) taskIdAndResult;
            String taskId = taskIdAndResultPair.getKey();
            String result = taskIdAndResultPair.getValue();
            switch (KeeperException.Code.get(rc)) {
                case CONNECTIONLOSS:
                    createStatusNodeAndDeleteTask(taskId, result);
                    break;
                case OK:
                    OpResult.CreateResult statusNodeCreateResult = (OpResult.CreateResult) opResults.get(0);
                    OpResult.DeleteResult taskNodeDeleteResult = (OpResult.DeleteResult) opResults.get(1);
                    LOG.info("create status node and delete task done: {}, {}", statusNodeCreateResult.getPath(), taskNodeDeleteResult.getType());
                    break;
                case NODEEXISTS:
                    LOG.warn("status node already created, {}", path);
                    break;
                default:
                    LOG.error("Something went wrong: " + KeeperException.create(KeeperException.Code.get(rc), path));
            }
        }
    };


    public static void main(String args[]) throws Exception {
        MyWorker myWorker = new MyWorker("127.0.0.1:2181");
        myWorker.startZK();

        while (!myWorker.connected) {
            Thread.sleep(100);
        }
        // 创建任务分配节点 和 worker节点
        myWorker.createWorkerAndAssignNode();
        // 获取分配的任务
        myWorker.getAssignedTasks();

        while (!myWorker.expired) {
            Thread.sleep(100);
        }

    }
}
